#!/usr/bin/env python
# -*- coding: utf-8 -*-


import re
import pandas.json as _json

from netCDF4 import Dataset
from metlib.kits import parse_slice_str
from metlib.datetime import parse_datetime, parse_timedelta

# ccmp_transporter = FileTransporter(remote_location='s3c://dxwind-windres-datasets', retry=2, retry_interval=1)
# ccmp_fc_manager = FileCacheManager('ccmp', ccmp_transporter)
# ccmp2d_fc_manager = FileCacheManager(bucket='ccmp2d',
#                                       transporter=ccmp_transporter)
# ccmp_kmz_manager = FileCacheManager('ccmp_kmz_img_tmp')


ccmp_py_varnames = ['mean', 'rose', 'dist']
ccmp_pt_varnames = ['ws10m', 'wd10m']
ccmp_pm_varnames = ['ws10m']
ccmp_ry_varnames = ['ws10m']
ccmp_years = [str(year) for year in range(1987, 2012)]
ccmp_zh_dataset_d = {
    'ccmp': u"CCMP海洋风场数据集"
}

ccmp_zh_subset_d = {
    "PT": u"时间序列",
    "PY": u"年统计值",
    "PM": u"月统计值",
    # "PD": u"日变化",
    "RY": u"年统计值",
}

ccmp_zh_varname_d = {
    "ws10m": u"风速(10m)",
    "wd10m": u"风向(10m)",
    "mean": u"汇总统计",
    "rose": u"风玫瑰",
    "windrose": u"风玫瑰",
    "dist": u"风速分布",
}

ccmp_units_d = {
    "ws10m": u"m/s",
    "wd10m": u"°",
    "summary": u"",
    "rose": u"‰",
    "windrose": u"‰",
    "dist": u"%",
}

ccmp_suggest_range_d = {
    "ws10m": (3.0, 10.0),
    "wd10m": (0.0, 360.0)
}

ccmp_cmap_d = {
    'ws10m': 'jet',
}

# ccmp海洋数据集的经纬度范围和分辨率
ccmp_lonlat_range = {
    'min_lon': 0.125,
    'max_lon': 359.875,
    'nx': 1440,
    'min_lat': -78.375,
    'max_lat': 78.375,
    'ny': 628,
    'lon_delta': 0.25,
    'lat_delta': 0.25,
    'years': ['1987', '1988', '1989', '1990', '1991', '1992', '1993', '1994',
              '1995', '1996', '1997', '1998', '1999', '2000', '2001', '2002',
              '2003', '2004', '2005', '2006', '2007', '2008', '2009', '2010', '2011']
}

ccmp_min_lon = ccmp_lonlat_range.get('min_lon')
ccmp_max_lon = ccmp_lonlat_range.get('max_lon')
ccmp_min_lat = ccmp_lonlat_range.get('min_lat')
ccmp_max_lat = ccmp_lonlat_range.get('max_lat')
ccmp_lon_delta = ccmp_lonlat_range.get('lon_delta')
ccmp_lat_delta = ccmp_lonlat_range.get('lat_delta')


def lonlat_to_ccmp_ij(lon, lat, info=ccmp_lonlat_range):
    """
    经纬度转为CCMP海洋数据集的格点坐标(ij)。
    """
    # lonlat limit
    lonlat_limit = info.get('limit')
    if lonlat_limit is not None:
        limit_min_lon = lonlat_limit.get('min_lon')
        limit_min_lat = lonlat_limit.get('min_lat')
        limit_max_lon = lonlat_limit.get('max_lon')
        limit_max_lat = lonlat_limit.get('max_lat')
        if lon > limit_max_lon or lon < limit_min_lon\
            or lat > limit_max_lat or lat < limit_min_lat:
            return

    min_lon = info.get('min_lon')
    min_lat = info.get('min_lat')
    lon_delta = info.get('lon_delta')
    lat_delta = info.get('lat_delta')
    i = int(round((lon - min_lon) / lon_delta))
    j = int(round((lat - min_lat) / lat_delta))
    # i, j在给定的格点范围内，返回(i, j), 不在返回None
    nx1 = 0
    ny1 = 0
    nx2 = info.get('nx')
    ny2 = info.get('ny')
    ij = None
    if nx1 <= i < nx2 and ny1 <= j < ny2:
        ij = (i, j)
    return ij


def latlon_to_ccmp_latlon(lat, lon, info=ccmp_lonlat_range):
    """
    经纬度转(latlon)为CCMP海洋数据集的格点经纬度坐标(latlon)。
    """
    ij = lonlat_to_ccmp_ij(lon, lat, info)
    lon_delta = info.get('lon_delta')
    lat_delta = info.get('lat_delta')
    if ij is not None:
        ix, jy = ij
        min_lon = info.get('min_lon')
        min_lat = info.get('min_lat')
        grid_lon = min_lon + ix * lon_delta
        grid_lat = min_lat + jy * lat_delta
        # ccmp_latlon = (grid_lat, grid_lon)
        ccmp_latlon = '%.3f_%.3f' % (grid_lat, grid_lon)
    else:
        ccmp_latlon = None
    return ccmp_latlon


def get_ccmp_rect(lon1, lat1, lon2, lat2, info):
    min_lon = min(lon1, lon2)
    max_lon = max(lon1, lon2)
    min_lat = min(lat1, lat2)
    max_lat = max(lat1, lat2)

    # CCMP海洋风场数据集的经纬度范围
    min_ccmp_lon = info.get('min_lon')
    max_ccmp_lon = info.get('max_lon')
    min_ccmp_lat = info.get('min_lat')
    max_ccmp_lat = info.get('max_lat')

    # lonlat limit
    lonlat_limit = info.get('limit')
    if lonlat_limit is not None:
        limit_min_lon = lonlat_limit.get('min_lon')
        limit_min_lat = lonlat_limit.get('min_lat')
        limit_max_lon = lonlat_limit.get('max_lon')
        limit_max_lat = lonlat_limit.get('max_lat')
        min_ccmp_lon = limit_min_lon if limit_min_lon > min_ccmp_lon else min_ccmp_lon
        max_ccmp_lon = limit_max_lon if limit_max_lon < max_ccmp_lon else max_ccmp_lon
        min_ccmp_lat = limit_min_lat if limit_min_lat > min_ccmp_lat else min_ccmp_lat
        max_ccmp_lat = limit_max_lat if limit_max_lat < max_ccmp_lat else max_ccmp_lat

    # 超出CCMP海洋风场范围
    if min_lon > max_ccmp_lon or max_lon < min_ccmp_lon \
            or min_lat > max_ccmp_lat or max_lat < min_ccmp_lat:
        return

    # 在CCMP海洋风场范围内
    lon_start = min_lon if min_lon > min_ccmp_lon else min_ccmp_lon
    lon_end = max_lon if max_lon < max_ccmp_lon else max_ccmp_lon
    lat_start = min_lat if min_lat > min_ccmp_lat else min_ccmp_lat
    lat_end = max_lat if max_lat < max_ccmp_lat else max_ccmp_lat

    min_grid_lat, min_grid_lon = latlon_to_ccmp_latlon(lat_start, lon_start, info).split('_')
    max_grid_lat, max_grid_lon = latlon_to_ccmp_latlon(lat_end, lon_end, info).split('_')
    return {
        'lat1': float(min_grid_lat),
        'lon1': float(min_grid_lon),
        'lat2': float(max_grid_lat),
        'lon2': float(max_grid_lon),
        'jy_ix': '%s:%s_%s:%s' % (min_grid_lat, max_grid_lat, min_grid_lon, max_grid_lon)
    }


def parse_ccmp_uri(uri):
    reg = r'(?P<dataset>[^/]+)/(?P<subset>[^/]+)/(?P<varname>[^/]+)/(?P<time>[^-+Z/]+)(Z?)(?P<timezone>[-+0-9]*)/(?P<jy_ix>[^/]+)$'
    m = re.match(reg, uri)
    if not m:
        raise ValueError('%s is not a proper ccmp uri' % uri)
    result = {field: m.group(field) for field in ['dataset', 'subset', 'varname', 'time', 'jy_ix']}

    time_str = result.get('time')
    timezone_str = m.group('timezone')
    try:
        timezone = int(timezone_str)
    except ValueError:
        timezone = 0

    if ',' in time_str:
        times = time_str.split(',')
        result['times'] = times
        result['timetype'] = 'year_list'
    elif ':' in time_str:
        begdt, enddt, tdelta = parse_slice_str(time_str, default_step='1h')
        begdt = parse_datetime(begdt)
        enddt = parse_datetime(enddt)
        result['begdt'] = begdt
        result['enddt'] = enddt
        result['tdelta'] = tdelta
        result['timetype'] = 'dtrange'
    elif len(m.group('time')) == 4:  # year
        begdt = parse_datetime(m.group('time')+'0101')
        enddt = begdt + parse_timedelta('1Y')
        result['begdt'] = begdt
        result['enddt'] = enddt
        result['tdelta'] = '1h'
        result['timetype'] = 'year'
    else:
        raise ValueError('%s is not a proper ccmp uri' % uri)

    result['timezone'] = timezone
    lat, lon = result['jy_ix'].split('_')
    if ':' in lat:
        lat1, lat2 = lat.split(':')
        lon1, lon2 = lon.split(':')
        lat1 = float(lat1); lat2 = float(lat2)
        lon1 = float(lon1); lon2 = float(lon2)
        mlat1, mlon1 = latlon_to_ccmp_latlon(lat1, lon1).split('_')
        mlat2, mlon2 = latlon_to_ccmp_latlon(lat2, lon2).split('_')
        jy_ix = '%s:%s_%s:%s' % (mlat1, mlat2, mlon1, mlon2)
    else:
        jy_ix = latlon_to_ccmp_latlon(float(lat), float(lon))
    result['jy_ix'] = jy_ix
    return result


class CCMP_PM_Handler(object):

    pm_desc_pattern = u"CCMP海洋风场数据集-月统计值\n变量: 平均值(单位: ws10m: m/s)\n格点经纬度:%.3f_%.3f\n年份:%s"

    def __init__(self, pm_file):
        self._pm = Dataset(pm_file, mode='r')
        self.lon = self._pm.variables['lon'][:]
        self.lat = self._pm.variables['lat'][:]
        self.tag = self._pm.variables['tag'][:]
        self.ws10m = self._pm.variables['ws10m']
        self.std = self._pm.variables['std']
        self.count = self._pm.variables['count'][:]
        self.subset_type = "monthlymean"


    def get_pm_common(self, grid_lon, grid_lat):
        return {
            "info": {
                "subset": "PM",
                "grid_lon": grid_lon,
                "grid_lat": grid_lat,
                "zh_dataset": "CCMP海洋风场数据集",
                "dataset": "ccmp",
                "time": "%s" % self.year,
                "jy_ix": "%.3f_%.3f" % (grid_lat, grid_lon),
                "zh_subset": "月统计值",
                "resolution": "0.25°"
            }
        }

    def get_pm_contents(self, grid_lon, grid_lat, variable='ws10m'):
        ij = lonlat_to_ccmp_ij(grid_lon, grid_lat)
        if ij is None:
            return
        ix, jy = ij
        vn_d = {
            "info": {
                "varname": variable,
                "type": self.subset_type,
                "zh_varname": ccmp_zh_varname_d.get(variable)
            },
            "data": {
                "std": self.std[:, jy, ix],
                "count": self.count,
                "name": variable,
                "zh_name": ccmp_zh_varname_d.get(variable),
                "values": self.ws10m[:, jy, ix],
                "units": ccmp_units_d.get(variable)
            },
            "uri": self.uri(grid_lon, grid_lat, variable)
        }
        return {
            variable: vn_d
        }

    def get_pm_data(self, grid_lon, grid_lat, variable='ws10m'):
        pm_data = {
            "desc": self.pm_desc_pattern % (grid_lat, grid_lon, self.year),
            "uri": self.uri(grid_lon, grid_lat, variable),
            "common": self.get_pm_common(grid_lon, grid_lat),
            "contents": self.get_pm_contents(grid_lon, grid_lat, variable),
            "tags": {
                variable: self.uri(grid_lon, grid_lat, variable)
            }
        }
        return pm_data

    def ccmp2json(self, dest_file, grid_lon, grid_lat, variable='ws10m'):
        data = self.get_pm_data(grid_lon, grid_lat, variable)
        with open(dest_file, mode='w') as fp:
            fp.write(_json.dumps(data, double_precision=4))

    def uri(self, grid_lon, grid_lat, variable):
        return "ccmp/PM/%s/%s/%.3f_%.3f" % (variable, self.year, grid_lat, grid_lon)

    @property
    def year(self):
        return self.tag[0][:4]

    def close(self):
        self._pm.close()


class CCMP_PY_Handler(object):
    py_mean_sub_varnames = ['ws10m']
    def __init__(self, py_file):
        self._py = Dataset(py_file, mode='r')
        self.tag = self._py.variables['tag'][:]
        self.ws10m = self._py.variables['ws10m']
        self.std = self._py.variables['std']
        self.count = self._py.variables['count'][:]

    def get_py_common(self, grid_lon, grid_lat):
        return {
            "info": {
                "subset": "PY",
                "grid_lon": grid_lon,
                "grid_lat": grid_lat,
                "zh_dataset": "CCMP海洋风场数据集",
                "dataset": "ccmp",
                "time": "%s" % self.year,
                "jy_ix": "%.3f_%.3f" % (grid_lat, grid_lon),
                "zh_subset": "年统计值",
                "resolution": "0.25°"
            }
        }

    def get_py_data(self, grid_lon, grid_lat):
        ij = lonlat_to_ccmp_ij(grid_lon, grid_lat)
        if ij is None:
            return
        ix, jy = ij
        data_info = {
            'sub_varnames': self.py_mean_sub_varnames,
            'type': 'summary'
        }
        data = {}
        for vn in self.py_mean_sub_varnames:
            data[vn] = {
                'name': vn,
                'values': self.ws10m[0, jy, ix],
                'std': self.std[0, jy, ix],
                'count': self.count[0],
                'units': ccmp_units_d.get(vn),
                'zh_name': ccmp_zh_varname_d.get(vn),
            }
        uri = u'ccmp/PY/mean/%s/%.3f_%.3f' % (self.year, grid_lat, grid_lon)
        return {
            'uri': uri,
            'common': self.get_py_common(grid_lon, grid_lat),
            'contents': {
                'data': {
                    'info': data_info,
                    'data': data,
                    'coords': {
                        'varname': 'mean'
                    }
                },
            },
            'tags': {
                'data': uri
            }
        }

    @property
    def year(self):
        return self.tag[0]

    def close(self):
        self._py.close()

